#include "DWI.h"

#include <teem/ten.h>

CDWI::CDWI()
{
	m_pDWI = NULL;
	m_pS0 = NULL;
	m_w = m_h = m_d = m_n = m_m = 0;
	m_pDiffusionGradients = NULL;
	m_bValue = 0;
}

CDWI::~CDWI()
{
	SafeDeleteArray(m_pDWI);
	SafeDeleteArray(m_pS0);
	m_w = m_h = m_d = m_n = m_m = 0;
	SafeDeleteArray(m_pDiffusionGradients);
	m_bValue = 0;
}

int CDWI::OpenDWIFile(const char *pathname, FILE *fp)
{
	if (pathname == NULL)
		return -1;

	fprintf(fp, "Reading DWI file %s ...", pathname);

	Nrrd *dwi = nrrdNew();
	if (nrrdLoad(dwi, pathname, NULL)) {
		fprintf(fp, "cannot load file %s:\n%s\n", pathname, biffGetDone(NRRD));
		return -2;			// cannot load file
	}

	// get the size of the DWI file
	m_n = (uint)(dwi->axis[0].size);
	m_w = (uint)(dwi->axis[1].size);
	m_h = (uint)(dwi->axis[2].size);
	m_d = (uint)(dwi->axis[3].size);

	// copy the DWI data
	m_pDWI = new float[m_w*m_h*m_d*m_n];
	memcpy(m_pDWI, dwi->data, sizeof(float)*m_w*m_h*m_d*m_n);

	// parse the DWI information
	Nrrd *gradKVP = NULL;
	Nrrd *bmatKVP = NULL;
	double bKVP;
	unsigned int *skip;
	unsigned int skipNum;
	char *err;
	airArray *mop;
	mop = airMopNew();
	if (tenDWMRIKeyValueParse(&gradKVP, &bmatKVP, &bKVP, &skip, &skipNum, dwi)) {
		airMopAdd(mop, err=biffGetDone(TEN), airFree, airMopAlways);
		fprintf(stderr, "trouble parsing DWI info:\n%s\n", err);
		airMopError(mop);
		return -3;
	}
	airMopOkay(mop);

	// B-value
	m_bValue = (float)bKVP;

	// diffusion gradients
	m_pDiffusionGradients = new CVector3F[m_n];
	for (uint i = 0; i < m_n; ++i) {
		if (gradKVP->type == nrrdTypeFloat) {
			m_pDiffusionGradients[i].m_x = ((float *)(gradKVP->data))[i*3+0];
			m_pDiffusionGradients[i].m_y = ((float *)(gradKVP->data))[i*3+1];
			m_pDiffusionGradients[i].m_z = ((float *)(gradKVP->data))[i*3+2];
		} else if (gradKVP->type == nrrdTypeDouble) {
			m_pDiffusionGradients[i].m_x = (float)((double *)(gradKVP->data))[i*3+0];
			m_pDiffusionGradients[i].m_y = (float)((double *)(gradKVP->data))[i*3+1];
			m_pDiffusionGradients[i].m_z = (float)((double *)(gradKVP->data))[i*3+2];
		}
	}

	// free data
	if (skipNum != 0)
		SafeDeleteArray(skip);
	nrrdNuke(gradKVP);
	nrrdNuke(bmatKVP);
	nrrdNuke(dwi);

	// compute S0
	m_m = 0;
	m_pS0 = new float[m_w*m_h*m_d];
	memset(m_pS0, 0, sizeof(float)*m_w*m_h*m_d);
	for (uint k = 0; k < m_n; ++k) {
		if (m_pDiffusionGradients[k].Length() < 0.01f) {
			// zero
			m_m++;
			for (uint i = 0; i < m_w*m_h*m_d; ++i) {
				m_pS0[i] += m_pDWI[i*m_n+k];
			}
		}
	}
	for (uint i = 0; i < m_w*m_h*m_d; ++i) {
		m_pS0[i] = m_pS0[i] / (float)m_m;
	}

	fprintf(fp, "Success.\n");
	fprintf(fp, "Width: %d, Height: %d, Depth: %d, gradients: %d\n", 
		m_w, m_h, m_d, m_n);

	return 0;
}

int CDWI::SaveDWIFile(const char *pathname, FILE *fp)
{
	if (pathname == NULL || m_pDWI == NULL)
		return -1;

	Nrrd *dst = nrrdNew();
	size_t ss[4];
	ss[0] = m_n;
	ss[1] = m_w;
	ss[2] = m_h;
	ss[3] = m_d;

	nrrdAlloc_nva(dst, nrrdTypeFloat, 4, ss);
	float *data = (float *)(dst->data);
	uint size = m_n * m_w * m_h * m_d;
	memcpy(data, m_pDWI, sizeof(float)*size);

	nrrdAxisInfoSet_va(dst, nrrdAxisInfoSpacing, 1.0, 1.0, 1.0, 1.0);
	nrrdAxisInfoSet_va(dst, nrrdAxisInfoKind, nrrdKindList, nrrdKindSpace, nrrdKindSpace, nrrdKindSpace);
	nrrdAxisInfoSet_va(dst, nrrdAxisInfoUnits, "", "mm", "mm", "mm");

	if (nrrdSave(pathname, dst, NULL)) {
		fprintf(fp, "cannot save file %s:\n%s\n", pathname, biffGetDone(NRRD));
		return -1;
	}

	nrrdNuke(dst);

	return 0;
}

int CDWI::CreateDWI(const uint n, const uint w, const uint h, const uint d, const float *pDWI)
{
	if (m_pDWI != NULL)
		return -1;

	uint size = n * w * h * d;
	if (size == 0)
		return -2;

	m_pDWI = new float[size];
	memset(m_pDWI, 0, sizeof(float)*size);
	m_n = n;
	m_w = w;
	m_h = h;
	m_d = d;

	if (pDWI != NULL) {
		memcpy(m_pDWI, pDWI, sizeof(float)*size);
	}

	return 0;
}

void CDWI::GetDWIFast(const uint x, const uint y, const uint z, float *dwi) const
{
	uint index = ((z * m_h + y) * m_w + x) * m_n;
	for (uint i = 0; i < m_n; ++i) {
		dwi[i] = m_pDWI[index+i];
	}
}


void CDWI::GetDWIFast(const CPoint3I &p, float *dwi) const
{
	GetDWIFast(p.x, p.y, p.z, dwi);
}

void CDWI::GetDWIFast(const CPoint3I *p, float *dwi) const
{
	GetDWIFast(p->x, p->y, p->z, dwi);
}

void CDWI::GetDWIFast(const uint index, float *dwi) const
{
	for (uint i = 0; i < m_n; ++i) {
		dwi[i] = m_pDWI[index*m_n+i];
	}
}

void CDWI::GetNS0(const uint index, float *dwi) const
{
	for (uint i = m_m; i < m_n; ++i) {
		dwi[i-m_m] = m_pDWI[index*m_n+i];
	}
}

void CDWI::GetDWI(const float x, const float y, const float z, float *dwi) const
{
	memset(dwi, 0, sizeof(float)*m_n);

	if (x < 0.0f || x > (float)(m_w-1) ||
		y < 0.0f || y > (float)(m_h-1) ||
		z < 0.0f || z > (float)(m_d-1)) {
			return;
	}

	uint x1 = (uint)floor(x);
	uint y1 = (uint)floor(y);
	uint z1 = (uint)floor(z);

	uint x2 = (uint)ceil(x);
	uint y2 = (uint)ceil(y);
	uint z2 = (uint)ceil(z);

	float a = x - (float)x1;
	float b = y - (float)y1;
	float c = z - (float)z1;

	uint i1 = ((z1 * m_h + y1) * m_w + x1) * m_n;
	uint i2 = ((z1 * m_h + y1) * m_w + x2) * m_n;
	uint i3 = ((z1 * m_h + y2) * m_w + x2) * m_n;
	uint i4 = ((z1 * m_h + y2) * m_w + x1) * m_n;
										     
	uint i5 = ((z2 * m_h + y1) * m_w + x1) * m_n;
	uint i6 = ((z2 * m_h + y1) * m_w + x2) * m_n;
	uint i7 = ((z2 * m_h + y2) * m_w + x2) * m_n;
	uint i8 = ((z2 * m_h + y2) * m_w + x1) * m_n;

	for (uint i = 0; i < m_n; ++i) {
		float v1 = m_pDWI[i1+i];
		float v2 = m_pDWI[i2+i];
		float v3 = m_pDWI[i3+i];
		float v4 = m_pDWI[i4+i];
		float v5 = m_pDWI[i5+i];
		float v6 = m_pDWI[i6+i];
		float v7 = m_pDWI[i7+i];
		float v8 = m_pDWI[i8+i];

		float temp = v1 + a * (v2 - v1) + b * (v4 - v1) + c * (v5 - v1) + 
			a * b * (v1 - v2 + v3 - v4) + 
			a * c * (v1 - v2 + v6 - v5) + 
			b * c * (v1 - v4 + v8 - v5) + 
			a * b * c * (-v1 + v2 - v3 + v4 + v5- v6 + v7 - v8);
		dwi[i] = temp;
	}
}

void CDWI::GetDWI(const CPoint3F &p, float *dti) const
{
	return GetDWI(p.m_x, p.m_y, p.m_z, dti);
}

void CDWI::GetDWI(const CPoint3F *p, float *dti) const
{
	return GetDWI(p->m_x, p->m_y, p->m_z, dti);
}

void CDWI::SetDWI(const uint x, const uint y, const uint z, const float *dwi)
{
	uint index = ((z * m_h + y) * m_w + x) * m_n;
	for (uint i = 0; i < m_n; ++i) {
		m_pDWI[index+i] = dwi[i];
	}
}

float CDWI::GetS0(const CPoint3F &p) const
{
	return GetS0(p.m_x, p.m_y, p.m_z);
}

float CDWI::GetS0(const float x, const float y, const float z) const
{
	if (x < 0.0f || x > (float)(m_w-1) ||
		y < 0.0f || y > (float)(m_h-1) ||
		z < 0.0f || z > (float)(m_d-1)) {
			return 0.0f;
	}

	uint x1 = (uint)floor(x);
	uint y1 = (uint)floor(y);
	uint z1 = (uint)floor(z);

	uint x2 = (uint)ceil(x);
	uint y2 = (uint)ceil(y);
	uint z2 = (uint)ceil(z);

	float a = x - (float)x1;
	float b = y - (float)y1;
	float c = z - (float)z1;

	uint i1 = (z1 * m_h + y1) * m_w + x1;
	uint i2 = (z1 * m_h + y1) * m_w + x2;
	uint i3 = (z1 * m_h + y2) * m_w + x2;
	uint i4 = (z1 * m_h + y2) * m_w + x1;
										 
	uint i5 = (z2 * m_h + y1) * m_w + x1;
	uint i6 = (z2 * m_h + y1) * m_w + x2;
	uint i7 = (z2 * m_h + y2) * m_w + x2;
	uint i8 = (z2 * m_h + y2) * m_w + x1;

	float v1 = m_pS0[i1];
	float v2 = m_pS0[i2];
	float v3 = m_pS0[i3];
	float v4 = m_pS0[i4];
	float v5 = m_pS0[i5];
	float v6 = m_pS0[i6];
	float v7 = m_pS0[i7];
	float v8 = m_pS0[i8];

	float res = v1 + a * (v2 - v1) + b * (v4 - v1) + c * (v5 - v1) + 
		a * b * (v1 - v2 + v3 - v4) + 
		a * c * (v1 - v2 + v6 - v5) + 
		b * c * (v1 - v4 + v8 - v5) + 
		a * b * c * (-v1 + v2 - v3 + v4 + v5- v6 + v7 - v8);

	return res;
}

void CDWI::GetNS0(const CPoint3F &p, float *dwi) const
{
	return GetNS0(p.m_x, p.m_y, p.m_z, dwi);
}

void CDWI::GetNS0(const float x, const float y, const float z, float *dwi) const
{
	memset(dwi, 0, sizeof(float)*m_m);

	if (x < 0.0f || x > (float)(m_w-1) ||
		y < 0.0f || y > (float)(m_h-1) ||
		z < 0.0f || z > (float)(m_d-1)) {
			return;
	}

	uint x1 = (uint)floor(x);
	uint y1 = (uint)floor(y);
	uint z1 = (uint)floor(z);

	uint x2 = (uint)ceil(x);
	uint y2 = (uint)ceil(y);
	uint z2 = (uint)ceil(z);

	float a = x - (float)x1;
	float b = y - (float)y1;
	float c = z - (float)z1;

	uint i1 = ((z1 * m_h + y1) * m_w + x1) * m_n;
	uint i2 = ((z1 * m_h + y1) * m_w + x2) * m_n;
	uint i3 = ((z1 * m_h + y2) * m_w + x2) * m_n;
	uint i4 = ((z1 * m_h + y2) * m_w + x1) * m_n;
										     
	uint i5 = ((z2 * m_h + y1) * m_w + x1) * m_n;
	uint i6 = ((z2 * m_h + y1) * m_w + x2) * m_n;
	uint i7 = ((z2 * m_h + y2) * m_w + x2) * m_n;
	uint i8 = ((z2 * m_h + y2) * m_w + x1) * m_n;

	for (uint i = m_m; i < m_n; ++i) {
		float v1 = m_pDWI[i1+i];
		float v2 = m_pDWI[i2+i];
		float v3 = m_pDWI[i3+i];
		float v4 = m_pDWI[i4+i];
		float v5 = m_pDWI[i5+i];
		float v6 = m_pDWI[i6+i];
		float v7 = m_pDWI[i7+i];
		float v8 = m_pDWI[i8+i];

		float temp = v1 + a * (v2 - v1) + b * (v4 - v1) + c * (v5 - v1) + 
			a * b * (v1 - v2 + v3 - v4) + 
			a * c * (v1 - v2 + v6 - v5) + 
			b * c * (v1 - v4 + v8 - v5) + 
			a * b * c * (-v1 + v2 - v3 + v4 + v5- v6 + v7 - v8);
		dwi[i-m_m] = temp;
	}
}